{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "v_eWsiiMfu2q",
      "metadata": {
        "id": "v_eWsiiMfu2q"
      },
      "source": [
        "<a target=\"_blank\" href=\"https://colab.research.google.com/github/cohere-ai/notebooks/blob/main/notebooks/llmu/RAG_with_Connectors.ipynb\">\n",
        "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
        "</a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f9e30089",
      "metadata": {
        "id": "f9e30089"
      },
      "source": [
        "# RAG with Connectors\n",
        "\n",
        "This notebook shows how to build a RAG-powered chatbot with Cohere's Chat endpoint using connectors. The chatbot can extract relevant information from external documents and produce verifiable, inline citations in its responses.\n",
        "\n",
        "Read the accompanying [article here](https://txt.cohere.com/rag-connectors/).\n",
        "\n",
        "Connectors are ways of connecting to data sources. These data sources could be internal documents, document databases, the broader internet, or any other source of context which can inform the replies generated by the model.\n",
        "\n",
        "We'll use the web search connector, a Cohere-managed connector that you can use without additional setup.\n",
        "\n",
        "The diagram below provides an overview of what we’ll build."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "Gg0u4eKVubYI",
      "metadata": {
        "id": "Gg0u4eKVubYI"
      },
      "source": [
        "![rag-workflow-3.png]()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5IdmW-I9NXpq",
      "metadata": {
        "id": "5IdmW-I9NXpq"
      },
      "source": [
        "# Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "r2jcKQ6iLefn",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "r2jcKQ6iLefn",
        "outputId": "acb51f35-43a2-4567-d8e6-913f00d57df6"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m117.2/117.2 kB\u001b[0m \u001b[31m1.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m75.6/75.6 kB\u001b[0m \u001b[31m3.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m77.8/77.8 kB\u001b[0m \u001b[31m4.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m58.3/58.3 kB\u001b[0m \u001b[31m4.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ],
      "source": [
        "! pip install cohere -q"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "id": "90f134ba",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        },
        "id": "90f134ba",
        "outputId": "f2236cef-f274-4100-dbcd-333b826f5ee8"
      },
      "outputs": [],
      "source": [
        "import uuid\n",
        "import cohere\n",
        "from cohere import ChatConnector\n",
        "from typing import List\n",
        "\n",
        "co = cohere.Client(\"COHERE_API_KEY\") # Get your API key here: https://dashboard.cohere.com/api-keys"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "id": "EavECqgqNJ8g",
      "metadata": {
        "cellView": "form",
        "id": "EavECqgqNJ8g"
      },
      "outputs": [],
      "source": [
        "#@title Enable text wrapping in Google Colab\n",
        "\n",
        "from IPython.display import HTML, display\n",
        "\n",
        "def set_css():\n",
        "  display(HTML('''\n",
        "  <style>\n",
        "    pre {\n",
        "        white-space: pre-wrap;\n",
        "    }\n",
        "  </style>\n",
        "  '''))\n",
        "get_ipython().events.register('pre_run_cell', set_css)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d319ece1",
      "metadata": {
        "id": "d319ece1"
      },
      "source": [
        "# Create a chatbot\n",
        "\n",
        "In connector mode, most of the implementation is taken care of by the endpoint, including deciding whether to retrieve information, generating queries, retrieving documents, chunking and reranking documents (post-retrieval), and generating the response. This greatly simplifies our code.\n",
        "\n",
        "The `Chatbot` class below handles the interaction between the user and chatbot.  We define the connector for the chatbot to use with the attribute `self.connectors`. In this notebook, we will use Cohere's `“web-search”` connector, which runs searches against a browser in safe mode.\n",
        "\n",
        "The run() method contains the logic for getting the user message, displaying the chatbot response with citations, along with a way for the user to end the conversation.\n",
        "\n",
        "Then, the chatbot responds to the user message.  We call `co.chat()` and supply a `connectors` parameter to make the chatbot component use connector mode.  All of the remaining implementation is taken care of by the endpoint, up to generating the response.\n",
        "\n",
        "We also pass the `conversation_id` parameter, which retains the interactions between the user and the chatbot in the same conversation thread. We enable the `stream` parameter so we can stream the chatbot response.\n",
        "\n",
        "We then print the chatbot's response.  In the case that the external information was used to generate a response, we also display documents and in-line citations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "id": "e52d521d",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        },
        "id": "e52d521d",
        "outputId": "b0f90f1c-17c8-46fa-d471-b11059767ede"
      },
      "outputs": [],
      "source": [
        "class Chatbot:\n",
        "    def __init__(self, connectors: List[str]):\n",
        "        \"\"\"\n",
        "        Initializes an instance of the Chatbot class.\n",
        "\n",
        "        \"\"\"\n",
        "        self.conversation_id = str(uuid.uuid4())\n",
        "        self.connectors = [ChatConnector(id=connector) for connector in connectors]\n",
        "\n",
        "    def run(self):\n",
        "        \"\"\"\n",
        "        Runs the chatbot application.\n",
        "\n",
        "        \"\"\"\n",
        "        while True:\n",
        "            # Get the user message\n",
        "            message = input(\"User: \")\n",
        "\n",
        "            # Typing \"quit\" ends the conversation\n",
        "            if message.lower() == \"quit\":\n",
        "                print(\"Ending chat.\")\n",
        "                break\n",
        "            else:                       # If using Google Colab, remove this line to avoid printing the same thing twice\n",
        "              print(f\"User: {message}\") # If using Google Colab, remove this line to avoid printing the same thing twice\n",
        "\n",
        "            # Generate response\n",
        "            response = co.chat_stream(\n",
        "                    message=message,\n",
        "                    model=\"command-r-plus\",\n",
        "                    conversation_id=self.conversation_id,\n",
        "                    connectors=self.connectors,\n",
        "            )\n",
        "\n",
        "            # Print the chatbot response, citations, and documents\n",
        "            print(\"\\nChatbot:\")\n",
        "            citations = []\n",
        "            cited_documents = []\n",
        "\n",
        "            # Display response\n",
        "            for event in response:\n",
        "                if event.event_type == \"text-generation\":\n",
        "                    print(event.text, end=\"\")\n",
        "                elif event.event_type == \"citation-generation\":\n",
        "                    citations.extend(event.citations)\n",
        "                elif event.event_type == \"stream-end\":\n",
        "                    cited_documents = event.response.documents\n",
        "\n",
        "            # Display citations and source documents\n",
        "            if citations:\n",
        "              print(\"\\n\\nCITATIONS:\")\n",
        "              for citation in citations:\n",
        "                print(citation)\n",
        "\n",
        "              print(\"\\nDOCUMENTS:\")\n",
        "              for document in cited_documents:\n",
        "                print({'id': document['id'],\n",
        "                      'snippet': document['snippet'][:400] + '...',\n",
        "                      'title': document['title'],\n",
        "                      'url': document['url']})\n",
        "\n",
        "            print(f\"\\n{'-'*100}\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "rJbwGPksPfh0",
      "metadata": {
        "id": "rJbwGPksPfh0"
      },
      "source": [
        "# Run the chatbot"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c755140c",
      "metadata": {
        "id": "c755140c"
      },
      "source": [
        "We can now run the chatbot.  For this, we create the instance of `Chatbot` using Cohere's managed web-search connector.  Then we run the chatbot by invoking the `run()` method.\n",
        "\n",
        "The format of each citation is:\n",
        "- `start`: The starting point of a span where one or more documents are referenced\n",
        "- `end`: The ending point of a span where one or more documents are referenced\n",
        "- `text`: The text representing this span\n",
        "- `document_ids`: The IDs of the documents being referenced (`doc_0` being the ID of the first document passed to the `documents` creating parameter in the endpoint call, and so on)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "id": "99e5005b",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 999
        },
        "id": "99e5005b",
        "outputId": "4609e72c-df6f-4c77-8132-cc0e73b80eee"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "User: What is Cohere's LLM University\n",
            "\n",
            "Chatbot:\n",
            "Cohere's LLM University (LLMU) is a set of comprehensive learning resources for anyone interested in natural language processing (NLP), from beginners to advanced learners. The curriculum covers everything from the basics of LLMs to the most advanced topics, including generative AI. The course is designed to give learners a solid foundation in NLP and help them develop their own applications.\n",
            "\n",
            "CITATIONS:\n",
            "start=24 end=30 text='(LLMU)' document_ids=['web-search_0', 'web-search_1']\n",
            "start=36 end=75 text='set of comprehensive learning resources' document_ids=['web-search_1']\n",
            "start=101 end=134 text='natural language processing (NLP)' document_ids=['web-search_0', 'web-search_1']\n",
            "start=141 end=172 text='beginners to advanced learners.' document_ids=['web-search_0', 'web-search_1']\n",
            "start=177 end=187 text='curriculum' document_ids=['web-search_0', 'web-search_1']\n",
            "start=215 end=229 text='basics of LLMs' document_ids=['web-search_0', 'web-search_1']\n",
            "start=237 end=283 text='most advanced topics, including generative AI.' document_ids=['web-search_1']\n",
            "start=326 end=349 text='solid foundation in NLP' document_ids=['web-search_0', 'web-search_1']\n",
            "start=364 end=395 text='develop their own applications.' document_ids=['web-search_0', 'web-search_1']\n",
            "\n",
            "DOCUMENTS:\n",
            "{'id': 'web-search_0', 'snippet': 'Guides and ConceptsAPI ReferenceRelease NotesApplication ExamplesLLMU\\n\\nCoralDashboardDocumentationPlaygroundCommunityLog In\\n\\nCoralDashboardDocumentationPlaygroundCommunityLog In\\n\\nWelcome to LLM University!\\n\\nWelcome to LLM University by Cohere!\\n\\nWe’re so happy that you’ve chosen to learn Natural Language Processing and Large Language Models with us.\\n\\nOur comprehensive curriculum aims to give you a ...', 'title': 'LLM University (LLMU) | Cohere', 'url': 'https://docs.cohere.com/docs/llmu'}\n",
            "{'id': 'web-search_1', 'snippet': \"Introducing LLM University — Your Go-To Learning Resource for NLP🎓\\n\\nDiscover our comprehensive NLP curriculum at LLM University. From the fundamentals of LLMs all the way to the most advanced topics, including generative AI\\n\\nWe're excited to announce the launch of LLM University (LLMU), a set of comprehensive learning resources for anyone interested in natural language processing (NLP), from begin...\", 'title': 'Introducing LLM University — Your Go-To Learning Resource for NLP🎓', 'url': 'https://txt.cohere.com/llm-university/'}\n",
            "\n",
            "----------------------------------------------------------------------------------------------------\n",
            "\n",
            "Ending chat.\n"
          ]
        }
      ],
      "source": [
        "# Define the connector\n",
        "connectors = [\"web-search\"]\n",
        "\n",
        "# Create an instance of the Chatbot class\n",
        "chatbot = Chatbot(connectors)\n",
        "\n",
        "# Run the chatbot\n",
        "chatbot.run()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.4"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}
